from urllib.request import urlopen
import os

import numpy as np
import cv2

from albumentations import (
    RandomBrightnessContrast,
    GaussianBlur,
    GaussNoise,
    Downscale,
    BboxParams,
    ShiftScaleRotate,
    HorizontalFlip,
    VerticalFlip,
    Rotate,
    Resize,
    CoarseDropout,
    RandomRotate90,
    RandomSizedBBoxSafeCrop,
    CenterCrop,
    RandomCrop,
    RandomSizedCrop,
    Crop,
    Compose,
    Transpose,
    OneOf
)

BOX_COLOR = (255, 0, 0)
TEXT_COLOR = (255, 255, 255)


def visualize_bbox(img, bbox, class_id, class_idx_to_name, color=BOX_COLOR, thickness=2):
    x_min, y_min, w, h = bbox
    x_min, x_max, y_min, y_max = int(x_min), int(x_min + w), int(y_min), int(y_min + h)
    cv2.rectangle(img, (x_min, y_min), (x_max, y_max), color=color, thickness=thickness)
    class_name = class_idx_to_name[class_id]
    ((text_width, text_height), _) = cv2.getTextSize(class_name, cv2.FONT_HERSHEY_SIMPLEX, 0.35, 1)
    cv2.rectangle(img, (x_min, y_min - int(1.3 * text_height)), (x_min + text_width, y_min), BOX_COLOR, -1)
    cv2.putText(img, class_name, (x_min, y_min - int(0.3 * text_height)), cv2.FONT_HERSHEY_SIMPLEX, 0.35,TEXT_COLOR, lineType=cv2.LINE_AA)
    return img


def visualize(annotations, category_id_to_name):
    def draw_caption(image, box, caption):
        b = np.array(box).astype(int)
        cv2.putText(image, caption, (b[0], b[1] - 10), cv2.FONT_HERSHEY_PLAIN, 1, (0, 0, 0), 2)
        cv2.putText(image, caption, (b[0], b[1] - 10), cv2.FONT_HERSHEY_PLAIN, 1, (255, 255, 255), 1)

    img = annotations['image'].copy()
    for idx,label_anno in enumerate(annotations['bboxes']):
        label_name = category_id_to_name[annotations['category_id'][idx]]
        print(label_anno)
        draw_caption(img, (int(label_anno[0]), int(label_anno[1]), int(label_anno[2]), int(label_anno[3])), label_name)
        cv2.rectangle(img, (int(label_anno[0]), int(label_anno[1])), (int(label_anno[2]), int(label_anno[3])), color=(0, 255, 255), thickness=1)
    cv2.imshow('img',img)
    cv2.waitKey(0)


def get_aug(aug,min_area=0., min_visibility=0.):
    return Compose(aug, bbox_params=BboxParams(format='pascal_voc', min_area=min_area,
                                               min_visibility=min_visibility, label_fields=['category_id']))

def resize(image,bboxes,category_id,size=(1024,512)):
    annotations = {'image':image,'bboxes':bboxes,'category_id':category_id}
    resize_fun = get_aug([Resize(size[0],size[1])])
    resized_data = resize_fun(**annotations)
    return resized_data['image'],resized_data['bboxes'],resized_data['category_id']


def aug_fun(image,bboxes,category_id,height,width):
    annotations = {'image': image, 'bboxes': bboxes, 'category_id': category_id}
    fun = get_aug([HorizontalFlip(p = 0.5),
                   # VerticalFlip(p = 0.2),
                   # Transpose(p = 0.2),
                   # RandomRotate90(p = 0.3),
                   Downscale(scale_min=0.6,scale_max=0.8,p = 0.5),
                   RandomSizedBBoxSafeCrop(p = 0.9,height=height,width=width),
                   ShiftScaleRotate(rotate_limit=20,p = 0.9),
                   RandomBrightnessContrast(p = 0.95),

                   CoarseDropout(max_holes=20,max_height=12,max_width=18,min_holes=4,min_height=5,min_width=8,p = 0.7),
                   OneOf([

                       GaussNoise(p = 0.7),
                       GaussianBlur(p = 0.7),
                   ])
                   ])
    return_data = fun(**annotations)
    return return_data['image'],return_data['bboxes'],return_data['category_id']

#
# aug_fun = get_aug([HorizontalFlip(p = 0.5),
#                    # VerticalFlip(p = 0.2),
#                    # Transpose(p = 0.2),
#                    # RandomRotate90(p = 0.3),
#                    Downscale(scale_min=0.6,scale_max=0.8,p = 0.5),
#                    RandomSizedBBoxSafeCrop(p = 0.9,height=1024,width=512),
#                    ShiftScaleRotate(rotate_limit=20,p = 0.9),
#                    RandomBrightnessContrast(p = 0.95),
#
#                    CoarseDropout(max_holes=20,max_height=12,max_width=18,min_holes=4,min_height=5,min_width=8,p = 0.7),
#                    OneOf([
#
#                        GaussNoise(p = 0.7),
#                        GaussianBlur(p = 0.7),
#                    ])
#                    ])
if __name__ == '__main__':
    image = cv2.imread('../data/train/img_0001885.jpg')
    print(type(image))
    annotations = {'image': image, 'bboxes': [[423, 168, 437, 179], [294, 210, 411, 291], [235, 306, 282, 394]],
                   'category_id': [1, 1, 1]}
    category_id_to_name = {1: 'pgps'}
    aug = get_aug([ShiftScaleRotate(p = 1),])
    augmented = aug(**annotations)

    visualize(annotations, category_id_to_name)
    visualize(augmented, category_id_to_name)